{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import random\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from PIL import Image\n",
    "import matplotlib.pylab as plt\n",
    "from loss import compute_loss\n",
    "from flownet import FlowNet\n",
    "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
    "from utils import split_lr_disp\n",
    "\n",
    "################################产生数据\n",
    "\n",
    "def trainGenerator(batch_size):\n",
    "\n",
    "    aug_dict = dict(horizontal_flip=True,\n",
    "                        fill_mode='nearest')\n",
    "    aug_dict = dict(horizontal_flip=True,\n",
    "                        fill_mode='nearest')\n",
    "\n",
    "    left_datagen = ImageDataGenerator(**aug_dict)\n",
    "    right_datagen = ImageDataGenerator(**aug_dict)\n",
    "    left_image = left_datagen.flow_from_directory(\n",
    "        \"F:/新建文件夹/Finputs1/frames_cleanpass/TRAIN/A\",\n",
    "        classes=None,\n",
    "        color_mode = \"rgb\",\n",
    "        target_size = (256, 512),\n",
    "        class_mode = None,\n",
    "        batch_size = batch_size, seed=1)\n",
    "    \n",
    "    right_image = right_datagen.flow_from_directory(\n",
    "        \"F:/新建文件夹/Finputs2/frames_cleanpass/TRAIN/A\",\n",
    "        classes=None,\n",
    "        color_mode = \"rgb\",\n",
    "        target_size = (256, 512),\n",
    "        class_mode = None,\n",
    "        batch_size = batch_size, seed=1)\n",
    "\n",
    "    train_generator = zip(left_image, right_image)\n",
    "    for (left,right) in train_generator:\n",
    "        img = left/255.\n",
    "        mask = right/255.        \n",
    "        yield (img,mask)\n",
    "        \n",
    "################################显示图片\n",
    "\n",
    "def show_single_image(img_arr):\n",
    "    plt.imshow(img_arr,cmap='binary')\n",
    "    plt.axis('on')\n",
    "    plt.show()\n",
    "\n",
    "################################超参\n",
    "    \n",
    "epochs = 1\n",
    "model = MonodepthNetwork()\n",
    "optimizer = tf.keras.optimizers.Adam()\n",
    "trainset = trainGenerator(batch_size=4)\n",
    "\n",
    "################################训练\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    for step in range(1):\n",
    "        with tf.GradientTape() as tape:\n",
    "            left_images, right_images = next(trainset)\n",
    "            print(left_images.shape,right_images.shape)\n",
    "            lr_disp = model(left_images, training=True)\n",
    "            image_loss, disp_gradient_loss, lr_loss = compute_loss(left_images, right_images, lr_disp)\n",
    "            total_loss = image_loss + disp_gradient_loss + lr_loss\n",
    "            gradients = tape.gradient(total_loss, model.trainable_variables)\n",
    "            optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "            print(\"EPOCH %2d STEP %3d total_loss %.6f image_loss %.6f disp_gradient_loss %.6f lr_loss %.6f\" %(\n",
    "                    epoch, step, total_loss.numpy(), image_loss.numpy(), disp_gradient_loss.numpy(), lr_loss.numpy()))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
